package tech.blueglacier.codec;

import com.google.common.primitives.Bytes;
import org.apache.commons.lang3.StringUtils;
import org.apache.james.mime4j.codec.Base64InputStream;
import org.apache.james.mime4j.codec.DecodeMonitor;
import org.apache.james.mime4j.codec.QuotedPrintableInputStream;
import org.apache.james.mime4j.io.InputStreams;
import org.apache.james.mime4j.util.ByteArrayBuffer;
import org.apache.james.mime4j.util.CharsetUtil;
import org.apache.james.mime4j.util.ContentUtil;
import tech.blueglacier.configuration.AppConfig;

import java.io.IOException;
import java.nio.charset.Charset;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Static methods for decoding strings, byte arrays and encoded words.
 * Enhancement the function of {@link org.apache.james.mime4j.codec.DecoderUtil}
 * which cannot decode the multiple encoded words which is not an integral number of octets.
 *
 * @author zhangxinchen
 * @since 2020/7/27
 */
public class EnhancedDecoderUtil {
    private static final Pattern PATTERN_ENCODED_WORD = Pattern.compile("=\\?(.+?)\\?(\\w)\\?(.*?)\\?=", Pattern.DOTALL);

    /**
     * Concatenate and decode a string containing Multiple encoded words as defined by RFC 2047
     * which are the same charset and encoding, non encoded word can be decoded too.
     * Encoded words have the form =?charset?enc?encoded-text?= where enc is either 'Q'
     * or 'q' for quoted-printable and 'B' or 'b' for base64. Using fallback
     * charset if charset in encoded words is invalid.
     *
     * @param body the string to decode
     * @return the decoded string.
     * @throws IllegalArgumentException only if the DecodeMonitor strategy throws it (Strict parsing)
     */
    public static String concatAndDecodeWords(String body) throws IllegalArgumentException {
        return concatAndDecodeWords(body, DecodeMonitor.SILENT);
    }

    /**
     * Concatenate and decode a string containing Multiple encoded words as defined by RFC 2047
     * which are the same charset and encoding, non encoded word can be decoded too.
     * Encoded words have the form =?charset?enc?encoded-text?= where enc is either 'Q'
     * or 'q' for quoted-printable and 'B' or 'b' for base64. Using fallback
     * charset if charset in encoded words is invalid.
     *
     * @param body     the string to decode
     * @param fallback the fallback Charset to be used.
     * @return the decoded string.
     * @throws IllegalArgumentException only if the DecodeMonitor strategy throws it (Strict parsing)
     */
    public static String concatAndDecodeWords(String body, Charset fallback) throws IllegalArgumentException {
        return concatAndDecodeWords(body, DecodeMonitor.SILENT, fallback);
    }

    /**
     * Concatenate and decode a string containing Multiple encoded words as defined by RFC 2047
     * which are the same charset and encoding, non encoded word can be decoded too.
     * Encoded words have the form =?charset?enc?encoded-text?= where enc is either 'Q'
     * or 'q' for quoted-printable and 'B' or 'b' for base64.
     *
     * @param body    the string to decode
     * @param monitor the DecodeMonitor to be used.
     * @return the decoded string.
     * @throws IllegalArgumentException only if the DecodeMonitor strategy throws it (Strict parsing)
     */
    public static String concatAndDecodeWords(String body, DecodeMonitor monitor) throws IllegalArgumentException {
        return concatAndDecodeWords(body, monitor, null);
    }

    /**
     * Concatenate and decode a string containing Multiple encoded words as defined by RFC 2047
     * which are the same charset and encoding, non encoded word can be decoded too.
     * Encoded words have the form =?charset?enc?encoded-text?= where enc is either 'Q'
     * or 'q' for quoted-printable and 'B' or 'b' for base64. Using fallback
     * charset if charset in encoded words is invalid.
     *
     * @param body     the string to decode
     * @param monitor  the DecodeMonitor to be used.
     * @param fallback the fallback Charset to be used.
     * @return the decoded string.
     * @throws IllegalArgumentException only if the DecodeMonitor strategy throws it (Strict parsing)
     */
    public static String concatAndDecodeWords(String body, DecodeMonitor monitor, Charset fallback) throws IllegalArgumentException {
        if (StringUtils.isEmpty(body)) {
            return body;
        }

        int tailIndex = 0;
        String lastMimeCharset = null;
        String lastEncoding = null;
        byte[] lastDecodedBytes = new byte[0];
        boolean lastMatchValid = false;
        StringBuilder sb = new StringBuilder();

        for (Matcher matcher = PATTERN_ENCODED_WORD.matcher(body); matcher.find(); ) {
            //mime4j默认将原始字节数组编码为US-ASCII，我们将非encoded-words字符串转码成原始字节数组，再探测
            String separator = reEncodingNonEncodedWords(body.substring(tailIndex, matcher.start()));
            String mimeCharset = AppConfig.getFallbackCharset(matcher.group(1));
            String encoding = matcher.group(2);
            String encodedText = matcher.group(3);
            byte[] decodedBytes = new byte[0];
            if (!StringUtils.isEmpty(encodedText)) {
                decodedBytes = tryDecodeEncodedWord(mimeCharset, encoding, encodedText, monitor, fallback);
            }

            if (decodedBytes == null) {
                sb.append(separator);
                sb.append(matcher.group(0));
            } else {
                if (isFirstEncodedWord(lastMimeCharset, lastEncoding)
                        || canConcat(lastMimeCharset, mimeCharset, lastEncoding, encoding, separator)) {
                    lastDecodedBytes = Bytes.concat(lastDecodedBytes, decodedBytes);
                } else {
                    sb.append(flush(lastDecodedBytes, lastMimeCharset, fallback));
                    lastDecodedBytes = decodedBytes;
                }

                if (!lastMatchValid || !CharsetUtil.isWhitespace(separator)) {
                    sb.append(separator);
                }
            }

            tailIndex = matcher.end();
            lastMimeCharset = mimeCharset;
            lastEncoding = encoding;
            lastMatchValid = (decodedBytes != null);
        }

        if (lastDecodedBytes.length != 0) {
            sb.append(flush(lastDecodedBytes, lastMimeCharset, fallback));
        }

        if (tailIndex == 0) {
            //兼容整段字符串均为非encoded-word
            return reEncodingNonEncodedWords(body);
        } else {
            //尾部为非encoded-word
            String tail = reEncodingNonEncodedWords(body.substring(tailIndex));
            sb.append(tail);
            return sb.toString();
        }
    }

    private static String reEncodingNonEncodedWords(String nonEncodedWords) {
        return CharsetDetector.autoDetectEncodingWithDefaultIfAbsent(ContentUtil.encode(nonEncodedWords).toByteArray());
    }

    // return null on error
    private static byte[] tryDecodeEncodedWord(final String mimeCharset,
                                               final String encoding,
                                               final String encodedText,
                                               final DecodeMonitor monitor,
                                               final Charset fallback) {
        Charset charset = CharsetUtil.lookup(mimeCharset);
        if (charset == null) {
            if (fallback == null) {
                monitor(monitor, mimeCharset, encoding, encodedText, "leaving word encoded",
                        "Mime charser '", mimeCharset, "' doesn't have a corresponding Java charset");
                return null;
            } else {
                charset = fallback;
            }
        }

        if (encodedText.length() == 0) {
            monitor(monitor, mimeCharset, encoding, encodedText, "leaving word encoded",
                    "Missing encoded text in encoded word");
            return null;
        }

        try {
            if (encoding.equalsIgnoreCase("Q")) {
                return decodeQ(encodedText, monitor);
            } else if (encoding.equalsIgnoreCase("B")) {
                return decodeB(encodedText, monitor);
            } else {
                monitor(monitor, mimeCharset, encoding, encodedText, "leaving word encoded",
                        "Warning: Unknown encoding in encoded word");
                return null;
            }
        } catch (RuntimeException e) {
            monitor(monitor, mimeCharset, encoding, encodedText, "leaving word encoded",
                    "Could not decode (", e.getMessage(), ") encoded word");
            return null;
        }
    }

    private static void monitor(DecodeMonitor monitor,
                                String mimeCharset,
                                String encoding,
                                String encodedText,
                                String dropDesc,
                                String... strings) throws IllegalArgumentException {
        if (monitor.isListening()) {
            String encodedWord = recombine(mimeCharset, encoding, encodedText);
            StringBuilder text = new StringBuilder();
            for (String str : strings) {
                text.append(str);
            }
            text.append(" (");
            text.append(encodedWord);
            text.append(")");
            String exceptionDesc = text.toString();
            if (monitor.warn(exceptionDesc, dropDesc))
                throw new IllegalArgumentException(text.toString());
        }
    }

    private static String recombine(final String mimeCharset,
                                    final String encoding,
                                    final String encodedText) {
        return "=?" + mimeCharset + "?" + encoding + "?" + encodedText + "?=";
    }

    /**
     * Decodes an encoded text encoded with the 'B' encoding (described in
     * RFC 2047) found in a header field body.
     *
     * @param encodedText the encoded text to decode.
     * @param monitor     {@link DecodeMonitor}
     * @return the decoded bytes.
     */
    private static byte[] decodeB(String encodedText, DecodeMonitor monitor) {
        encodedText = encodedText.replaceAll(" ", "+");
        return decodeBase64(encodedText, monitor);
    }

    /**
     * Decodes a string containing base64 encoded data.
     *
     * @param s       the string to decode.
     * @param monitor {@link DecodeMonitor}
     * @return the decoded bytes.
     */
    private static byte[] decodeBase64(String s, DecodeMonitor monitor) {
        try (Base64InputStream is = new Base64InputStream(InputStreams.createAscii(s), monitor)) {
            ByteArrayBuffer buf = new ByteArrayBuffer(s.length());
            int b;
            while ((b = is.read()) != -1) {
                buf.append(b);
            }
            return buf.toByteArray();
        } catch (IOException ex) {
            // This should never happen!
            throw new Error(ex);
        }
    }

    /**
     * Decodes a string containing quoted-printable encoded data.
     *
     * @param s the string to decode.
     * @return the decoded bytes.
     */
    private static byte[] decodeQuotedPrintable(String s, DecodeMonitor monitor) {
        try (QuotedPrintableInputStream is = new QuotedPrintableInputStream(InputStreams.createAscii(s), monitor)) {
            ByteArrayBuffer buf = new ByteArrayBuffer(s.length());
            int b;
            while ((b = is.read()) != -1) {
                buf.append(b);
            }
            return buf.toByteArray();
        } catch (IOException ex) {
            // This should never happen!
            throw new Error(ex);
        }
    }

    /**
     * Decodes an encoded text encoded with the 'Q' encoding (described in
     * RFC 2047) found in a header field body.
     *
     * @param encodedText the encoded text to decode.
     * @return the decoded bytes.
     */
    private static byte[] decodeQ(String encodedText, DecodeMonitor monitor) {
        return decodeQuotedPrintable(encodedText, monitor);
    }

    private static boolean isFirstEncodedWord(String lastMimeCharset, String lastEncoding) {
        return lastMimeCharset == null && lastEncoding == null;
    }

    private static boolean canConcat(String lastMimeCharset,
                                     String mimeCharset,
                                     String lastEncoding,
                                     String encoding,
                                     String separator) {
        return sameCharsetAndEncoding(lastMimeCharset, mimeCharset, lastEncoding, encoding) && CharsetUtil.isWhitespace(separator);
    }

    private static boolean sameCharsetAndEncoding(String lastMimeCharset,
                                                  String mimeCharset,
                                                  String lastEncoding,
                                                  String encoding) {
        return mimeCharset.equalsIgnoreCase(lastMimeCharset) && encoding.equalsIgnoreCase(lastEncoding);
    }

    private static String flush(byte[] decodedBytes, String mimeCharset, Charset fallbackCharset) throws IllegalArgumentException {
        Charset charset = CharsetUtil.lookup(mimeCharset);
        if (charset == null) {
            if (fallbackCharset == null) {
                throw new IllegalArgumentException("Unable to determine charset");
            } else {
                charset = fallbackCharset;
            }
        }

        return new String(decodedBytes, charset);
    }
}
